﻿using ConvNetSharp.Core.Serialization;
using System;
using System.Collections.Generic;
using System.Text;

namespace JustMathCore.ML
{
    public class ConvNetModel : IMLModel
    {
        ConvNetSharp.Core.Net<float> net;
        ConvNetSharp.Core.Training.Single.AdamTrainer trainer;
        public bool InitModel()
        {
            net = new ConvNetSharp.Core.Net<float>();
            trainer = new ConvNetSharp.Core.Training.Single.AdamTrainer(net);


            //输入层
            net.AddLayer(new ConvNetSharp.Core.Layers.InputLayer<float>(1, 1, 160));

            //全连接层
            net.AddLayer(new ConvNetSharp.Core.Layers.FullyConnLayer<float>(256));
            net.AddLayer(new ConvNetSharp.Core.Layers.FullyConnLayer<float>(512));
            net.AddLayer(new ConvNetSharp.Core.Layers.FullyConnLayer<float>(1024));
            net.AddLayer(new ConvNetSharp.Core.Layers.ReluLayer<float>());
            net.AddLayer(new ConvNetSharp.Core.Layers.FullyConnLayer<float>(512));

            //激活层
            net.AddLayer(new ConvNetSharp.Core.Layers.ReluLayer<float>());

            //规范分类
            net.AddLayer(new ConvNetSharp.Core.Layers.FullyConnLayer<float>(256));
            //输出层
            net.AddLayer(new ConvNetSharp.Core.Layers.RegressionLayer<float>());

            return true;
        }
        public bool LoadModel(string filename)
        {
            var jsonstr = System.IO.File.ReadAllText(filename);
            var _net = ConvNetSharp.Core.Serialization.SerializationExtensions.FromJson<float>(jsonstr);
            this.net = _net;
            trainer = new ConvNetSharp.Core.Training.Single.AdamTrainer(net);
            return true;
        }

        public void SaveModel(string filename)
        {
            var jsonstr = net.ToJson();
            System.IO.File.Delete(filename);
            System.IO.File.WriteAllText(filename, jsonstr);
        }


        [ThreadStatic]
        static float[] xbuf;
        [ThreadStatic]
        static float[] ybuf;
        static void ByteToFloat(byte[] data, float[] result)
        {
            int len = data.Length;
            if (len != 20 && len != 32) throw new Exception("must be 160 bit or 256 bit");

            for (var i = 0; i < len; i++)
            {
                var b = data[i];
                result[i * 8 + 0] = b & 0x01;
                result[i * 8 + 1] = b & 0x02;
                result[i * 8 + 2] = b & 0x04;
                result[i * 8 + 3] = b & 0x08;
                result[i * 8 + 4] = b & 0x10;
                result[i * 8 + 5] = b & 0x20;
                result[i * 8 + 6] = b & 0x40;
                result[i * 8 + 7] = b & 0x80;
            }
        }
        static void FloatToByte(float[] data, byte[] result)
        {
            int len = data.Length;
            if (len != 160 && len != 256) throw new Exception("must be 160 bit or 256 bit");
            for (var i = 0; i < len / 8; i++)
            {
                var b0 = data[i * 8 + 0] > 0.5f ? 1 : 0;
                var b1 = data[i * 8 + 1] > 0.5f ? 2 : 0;
                var b2 = data[i * 8 + 2] > 0.5f ? 4 : 0;
                var b3 = data[i * 8 + 3] > 0.5f ? 8 : 0;
                var b4 = data[i * 8 + 4] > 0.5f ? 0x10 : 0;
                var b5 = data[i * 8 + 5] > 0.5f ? 0x20 : 0;
                var b6 = data[i * 8 + 6] > 0.5f ? 0x40 : 0;
                var b7 = data[i * 8 + 7] > 0.5f ? 0x80 : 0;

                result[i] = (byte)(b0 | b1 | b2 | b3 | b4 | b5 | b6 | b7);
            }

        }



        public void Predict(byte[] x, byte[] yresult)
        {
            if (xbuf == null)
                xbuf = new float[160];
            ByteToFloat(x, xbuf);

            ConvNetSharp.Volume.Volume<float> vx = xbuf;

            float[] result = net.Forward(vx).ToArray();
            FloatToByte(result, yresult);
        }




        public void Train(byte[] x, byte[] y)
        {
            if (xbuf == null)
                xbuf = new float[160];
            ByteToFloat(x, xbuf);
            if (ybuf == null)
                ybuf = new float[256];
            ByteToFloat(y, ybuf);

            ConvNetSharp.Volume.Volume<float> vx = xbuf;
            ConvNetSharp.Volume.Volume<float> vy = ybuf;
            trainer.Train(vx, vy);
        }
        public float Loss
        {
            get
            {
                return trainer.Loss;
            }
        }

    }
}
